#!/usr/bin/env python3

import argparse
from glob import glob
import os
import re
import pendulum
import subprocess
from time import sleep
from shutil import copyfile
from netCDF4 import Dataset
import sys
sys.path.append(f'{os.path.dirname(os.path.realpath(__file__))}/../utils')
from utils import cli, check_files, run, submit_job

def copy_wrfda_output(dom_str, start_time_str, wrfda_work_dir):
	work_dir = wrfda_work_dir + '/' + dom_str
	if os.path.isfile(f'{work_dir}/wrfvar_output_{start_time_str}'):
		cli.notice(f'Use assimilated input for domain {dom_str}.')
	else:
		return False
	run(f'ln -sf {work_dir}/wrfvar_output_{start_time_str} wrfinput_{dom_str}')
	if dom_str == 'd01':
		run(f'ln -sf {work_dir}/wrfbdy_d01_{start_time_str}.lateral_updated wrfbdy_d01')
	return True

def run_wrf(work_root, wrf_root, config, args, wrfda_work_dir=None, tag=None):
	start_time = config['custom']['start_time']
	end_time = config['custom']['end_time']
	datetime_fmt = 'YYYY-MM-DD_HH:mm:ss'
	start_time_str = start_time.format(datetime_fmt)
	end_time_str = end_time.format(datetime_fmt)
	max_dom = config['domains']['max_dom']

	if not wrfda_work_dir:
		if tag != None:
			wrfda_work_dir = f'{work_root}/{tag}/wrfda'
		else:
			wrfda_work_dir = f'{work_root}/wrfda'
	elif not os.path.isdir(wrfda_work_dir):
		cli.error(f'run_wrf: {wrfda_work_dir} does not exist!')

	if tag != None:
		wrf_work_dir = f'{work_root}/{tag}/wrf'
	else:
		wrf_work_dir = f'{work_root}/wrf'
	if not os.path.isdir(wrf_work_dir): cli.error(f'run_wrf: {wrf_work_dir} does not exist!')
	old_dir = os.getcwd()
	os.chdir(wrf_work_dir)

	cli.stage(f'Run wrf.exe at {wrf_work_dir} ...')

	all_wrfda_ok = True
	for dom_idx in range(max_dom):
		dom_str = 'd' + str(dom_idx + 1).zfill(2)
		if not copy_wrfda_output(dom_str, start_time_str, wrfda_work_dir):
			all_wrfda_ok = False
			break
	if not all_wrfda_ok:
		cli.warning('Do not use data assimilation.')
		expected_files = ['wrfinput_d{:02d}_{}'.format(i + 1, start_time_str) for i in range(max_dom)]
		expected_files.append(f'wrfbdy_d01')
		if not check_files(expected_files):
			cli.error('real.exe wasn\'t executed successfully!')
		for i in range(max_dom):
			run('ln -sf wrfinput_d{0:02d}_{1} wrfinput_d{0:02d}'.format(i + 1, start_time_str))
		run(f'ln -sf wrfbdy_d01_{start_time_str} wrfbdy_d01')

	expected_files = []
	for i in range(max_dom):
		for time in pendulum.interval(start_time, end_time).range('minutes', config['time_control']['history_interval'][i]):
			pass
		expected_files.append(f'wrfout_d{i+1:02d}_{time.format(datetime_fmt)}')
	if not check_files(expected_files) or args.force:
		run('rm -f wrfout_*')
		run(f'ln -sf {wrf_root}/run/LANDUSE.TBL .')
		run(f'ln -sf {wrf_root}/run/ozone_plev.formatted .')
		run(f'ln -sf {wrf_root}/run/ozone_lat.formatted .')
		run(f'ln -sf {wrf_root}/run/ozone.formatted .')
		run(f'ln -sf {wrf_root}/run/RRTM_DATA_DBL RRTM_DATA')
		run(f'ln -sf {wrf_root}/run/RRTMG_LW_DATA .')
		run(f'ln -sf {wrf_root}/run/RRTMG_SW_DATA .')
		run(f'ln -sf {wrf_root}/run/VEGPARM.TBL .')
		run(f'ln -sf {wrf_root}/run/SOILPARM.TBL .')
		run(f'ln -sf {wrf_root}/run/GENPARM.TBL .')
		run(f'ln -sf {wrf_root}/run/CAMtr_volume_mixing_ratio .')
		submit_job(f'{wrf_root}/run/wrf.exe', args.np, config, args, wait=True, exclusive=True)
		if not check_files(expected_files):
			run(f'ls -l {wrf_work_dir}/wrfout_*')
			cli.warning(f'Failed! Check output {os.path.abspath(wrf_work_dir)}/rsl.error.0000.')
		cli.notice('Succeeded.')
	else:
		cli.notice('File wrfout_* already exist.')
	run(f'ls -l {wrf_work_dir}/wrfout_*')
	os.chdir(old_dir)
